'''
Created on 2020年1月30日
获取所有股票数据
@author: JM
'''
import pandas
import pandas as pd

from app_config import get_engine_ts, get_pro, get_engine
from sqlalchemy import text

engine = get_engine()


def read_data(index_code: str, date_str: str):
    sql = f"""
     select * from index_weight where index_code = '{index_code}' and trade_date = '{date_str}'
    """
    df = pd.read_sql_query(sql, engine)
    return df


def write_data(data_frame):
    res = data_frame.to_sql('index_weight', engine, index=False, if_exists='append', chunksize=5000)
    print(res)


def get_data(index_code: str, date_str: str) -> pd.DataFrame:
    data = get_pro().index_weight(index_code=index_code, start_date=date_str, end_date=date_str)
    return data


def read_get_data(index_code: str, date_str: str) -> pd.DataFrame:
    data: pd.DataFrame = read_data(index_code, date_str)
    if data.empty:
        data: pd.DataFrame = get_data(index_code, date_str)
        if data.empty:
            raise ValueError(f"index_weight is empty, index_code: {index_code}, date: {date_str}")
        write_data(data)
    return data


if __name__ == '__main__':
    read_get_data('399975.CSI', date_str='20250530')
    read_get_data('399975.CSI', date_str='20250630')
    read_get_data('399975.CSI', date_str='20250731')
    read_get_data('399975.CSI', date_str='20250829')
    read_get_data('399975.CSI', date_str='20250829')
    read_get_data('399975.CSI', date_str='20250930')
    read_get_data('399975.CSI', date_str='20251031')
    pass
